import cv2
import os
import sys
from tkinter import filedialog, Tk  # 添加文件对话框支持

# 在文件顶部添加Pillow的导入
from PIL import Image, ImageDraw, ImageFont
import numpy as np

class YOLOLabelTool:
    def __init__(self):
        self.original_image = None
        self.aspect_ratio = 1.0
        self.image = None
        self.annotations = []
        self.current_class = 0
        self.class_names = ["person", "car", "dog", "cat"]
        self.drawing = False
        self.ix, self.iy = -1, -1
        self.fx, self.fy = -1, -1
        self.language = "zh"  # 默认中文
        self.hover_index = -1  # 当前悬停的框索引
        self.inputting_label = False
        self.current_label_text = ""
        self.label_input_pos = (0, 0)
        # 添加中文字体路径初始化
        self.font_path = "C:/Windows/Fonts/simhei.ttf"  # 黑体字
        self.scale_factor = 1.0  # 初始缩放比例
        self.min_scale = 0.1    # 最小缩放比例
        self.max_scale = 10.0   # 最大缩放比例
        self.mouse_x, self.mouse_y = 0, 0  # 记录鼠标位置
        self.offset_x, self.offset_y = 0, 0  # 图像偏移量
        self.scale_x = 1.0  # 水平缩放比例
        self.scale_y = 1.0  # 垂直缩放比例
        
        # 添加教程文本
        self.tutorial_text_zh = """
YOLO标注工具使用教程:
1. 鼠标左键拖动 - 绘制标注框
2. 鼠标右键 - 删除当前悬停的标注框
3. 双击左键 - 编辑当前标注框的标签
4. 空格键 - 确认当前绘制的标注框
5. N键 - 切换到下一个标签类别
6. P键 - 切换到上一个标签类别
7. S键 - 保存标注到文件
8. C键 - 清除所有标注
9. L键 - 切换中英文
10. Q键 - 退出程序
11. 鼠标滚轮 - 缩放图像
"""
        self.tutorial_text_en = """
YOLO Label Tool Tutorial:
1. Left mouse drag - Draw bounding box
2. Right click - Delete hovered box
3. Double click - Edit box label
4. Space - Confirm current box
5. N - Next label class
6. P - Previous label class
7. S - Save annotations
8. C - Clear all annotations
9. L - Toggle language
10. Q - Quit
11. Mouse wheel - Zoom image
"""

    def load_image(self, image_path):
        # 检查路径是否存在
        if not os.path.exists(image_path):
            print(f"路径不存在: {image_path}")
            return False
            
        # 检查是否是文件(不是目录)
        if not os.path.isfile(image_path):
            print(f"请输入具体的图像文件路径，而不是目录: {image_path}")
            return False
            
        # 检查文件扩展名
        valid_extensions = ['.jpg', '.jpeg', '.png', '.bmp']
        if not any(image_path.lower().endswith(ext) for ext in valid_extensions):
            print(f"不支持的图像格式，请使用: {', '.join(valid_extensions)}")
            return False
            
        self.image = cv2.imread(image_path)
        if self.image is None:
            print(f"无法加载图像(文件可能损坏): {image_path}")
            return False
            
        # 设置默认窗口大小
        self.window_width = 800
        self.window_height = 600
        return True

    def run(self, image_path, output_path):
        if not self.load_image(image_path):
            return
        
        # 保存原始图像用于缩放
        self.original_image = self.image.copy()
        h, w = self.original_image.shape[:2]
        self.aspect_ratio = w / h
        
        cv2.namedWindow('YOLO Label Tool', cv2.WINDOW_NORMAL)
        cv2.setMouseCallback('YOLO Label Tool', self.draw_rectangle)
        cv2.resizeWindow('YOLO Label Tool', self.window_width, self.window_height)
        
        while True:
            # 获取当前窗口大小
            win_width = cv2.getWindowImageRect('YOLO Label Tool')[2]
            win_height = cv2.getWindowImageRect('YOLO Label Tool')[3]
            
            # 根据窗口大小调整缩放基准
            if win_width != self.window_width or win_height != self.window_height:
                self.window_width, self.window_height = win_width, win_height
                # 重置偏移量
                self.offset_x, self.offset_y = 0, 0
            
            # 计算缩放后的图像尺寸
            scaled_width = int(self.window_width * self.scale_factor)
            scaled_height = int(self.window_height * self.scale_factor)
            scaled_img = cv2.resize(self.original_image, (scaled_width, scaled_height))
            
            # 计算显示区域(考虑偏移量)
            x1 = int(max(0, min(self.offset_x, scaled_width - self.window_width)))
            y1 = int(max(0, min(self.offset_y, scaled_height - self.window_height)))
            x2 = int(min(x1 + self.window_width, scaled_width))
            y2 = int(min(y1 + self.window_height, scaled_height))
            
            # 裁剪图像
            self.image = scaled_img[y1:y2, x1:x2]
            
            # 如果裁剪后尺寸不匹配窗口(边缘情况)，填充黑色背景
            if self.image.shape[0] != self.window_height or self.image.shape[1] != self.window_width:
                pad_img = np.zeros((self.window_height, self.window_width, 3), dtype=np.uint8)
                pad_img[:self.image.shape[0], :self.image.shape[1]] = self.image
                self.image = pad_img
            
            img_copy = self.image.copy()
            
            # 计算缩放比例和偏移量
            self.scale_x = scaled_width / self.original_image.shape[1]
            self.scale_y = scaled_height / self.original_image.shape[0]
            offset_x = x1 / self.scale_x
            offset_y = y1 / self.scale_y

            # 绘制所有标注框（应用缩放和偏移）
            colors = [(0, 255, 0), (255, 0, 0), (0, 0, 255), (255, 255, 0)]
            for i, ann in enumerate(self.annotations):
                # 将原始坐标转换为当前显示坐标
                x1_disp = int((ann['x1'] - offset_x) * self.scale_x)  # 修改为self.scale_x
                y1_disp = int((ann['y1'] - offset_y) * self.scale_y)  # 修改为self.scale_y
                x2_disp = int((ann['x2'] - offset_x) * self.scale_x)  # 修改为self.scale_x
                y2_disp = int((ann['y2'] - offset_y) * self.scale_y)  # 修改为self.scale_y
                
                color = colors[i % len(colors)]
                thickness = 3 if i == self.hover_index else 2
                cv2.rectangle(img_copy, (x1_disp, y1_disp), (x2_disp, y2_disp), color, thickness)
            
            # 绘制当前正在画的框（应用缩放和偏移）
            if self.drawing and self.ix != -1 and self.fx != -1:
                ix_disp = int((self.ix - offset_x) * self.scale_x)  # 修改为self.scale_x
                iy_disp = int((self.iy - offset_y) * self.scale_y)  # 修改为self.scale_y
                fx_disp = int((self.fx - offset_x) * self.scale_x)  # 修改为self.scale_x
                fy_disp = int((self.fy - offset_y) * self.scale_y)  # 修改为self.scale_y
                cv2.rectangle(img_copy, (ix_disp, iy_disp), (fx_disp, fy_disp), (0, 0, 255), 2)
            
            # 将图像转换为PIL格式绘制文本
            img_pil = Image.fromarray(cv2.cvtColor(img_copy, cv2.COLOR_BGR2RGB))
            draw = ImageDraw.Draw(img_pil)
            
            # 使用PIL绘制文本（应用缩放和偏移）
            font = ImageFont.truetype(self.font_path, 20)
            small_font = ImageFont.truetype(self.font_path, 14)
            
            for i, ann in enumerate(self.annotations):
                # 转换标签坐标
                x1_disp = int((ann['x1'] - offset_x) * self.scale_x)  # 修改为self.scale_x
                y1_disp = int((ann['y1'] - offset_y) * self.scale_y)  # 修改为self.scale_y
                
                color = colors[i % len(colors)]
                label = f"{self.class_names[ann['class']]}"
                draw.text((x1_disp, y1_disp - 25), label, font=font, fill=color)
            
            # 绘制状态文本
            current_label = f"当前标签: {self.class_names[self.current_class]}" if self.language == "zh" else f"Current Label: {self.class_names[self.current_class]}"
            draw.text((10, 10), current_label, font=font, fill=(0, 0, 255))
            
            lang_hint = "按L切换中英文" if self.language == "zh" else "Press L to switch language"
            draw.text((10, 40), lang_hint, font=small_font, fill=(255, 255, 255))
            
            # 添加教程按钮提示
            tutorial_hint = "按H查看教程" if self.language == "zh" else "Press H for tutorial"
            draw.text((10, 70), tutorial_hint, font=small_font, fill=(255, 255, 255))
            
            # 转换回OpenCV格式并显示
            img_copy = cv2.cvtColor(np.array(img_pil), cv2.COLOR_RGB2BGR)
            cv2.imshow('YOLO Label Tool', img_copy)
            
            key = cv2.waitKey(1) & 0xFF
            if self.inputting_label:
                if key == 13:  # 回车键确认
                    if self.hover_index != -1:
                        self.class_names[self.annotations[self.hover_index]['class']] = self.current_label_text
                    self.inputting_label = False
                elif key == 8:  # 退格键
                    self.current_label_text = self.current_label_text[:-1]
                elif 32 <= key <= 126 or key > 127:  # 可打印字符(包括中文)
                    self.current_label_text += chr(key)
            else:
                if key == ord('q') or cv2.getWindowProperty('YOLO Label Tool', cv2.WND_PROP_VISIBLE) < 1:  # 退出或点击关闭按钮
                    break
                elif key == ord('s'):
                    if self.save_annotations(output_path):
                        print(f"标注已保存到: {output_path}" if self.language == "zh" else f"Annotations saved to: {output_path}")
                elif key == ord('c'):
                    self.annotations = []
                    print("已清除所有标注" if self.language == "zh" else "All annotations cleared")
                elif key == ord('n'):
                    self.current_class = (self.current_class + 1) % len(self.class_names)
                elif key == ord('p'):
                    self.current_class = (self.current_class - 1) % len(self.class_names)
                elif key == 32:  # 空格键
                    if self.drawing:
                        self.drawing = False
                        self.annotations.append({
                            'class': self.current_class,
                            'x1': min(self.ix, self.fx),
                            'y1': min(self.iy, self.fy),
                            'x2': max(self.ix, self.fx),
                            'y2': max(self.iy, self.fy)
                        })
                        self.ix, self.iy = -1, -1
                        self.fx, self.fy = -1, -1
                elif key == ord('l'):  # 切换语言
                    self.language = "en" if self.language == "zh" else "zh"
                elif key == ord('h'):  # 显示教程
                    tutorial_text = self.tutorial_text_zh if self.language == "zh" else self.tutorial_text_en
                    from tkinter import Toplevel, Text, Scrollbar, font as tkfont
                    root = Tk()
                    root.withdraw()
                    
                    # 创建自定义弹窗
                    help_window = Toplevel(root)
                    help_window.title("help")
                    help_window.geometry("800x500")
                    
                    # 设置字体
                    bold_font = tkfont.Font(family="宋体", size=10, weight="bold")
                    
                    # 添加滚动条
                    scrollbar = Scrollbar(help_window)
                    scrollbar.pack(side="right", fill="y")
                    
                    # 添加文本框
                    text = Text(help_window, 
                              wrap="word", 
                              yscrollcommand=scrollbar.set,
                              font=bold_font,
                              padx=10,
                              pady=10)
                    text.pack(expand=True, fill="both")
                    text.insert("end", tutorial_text)
                    text.config(state="disabled")  # 设为只读
                    
                    # 配置滚动条
                    scrollbar.config(command=text.yview)
                    
                    # 确保窗口关闭时释放资源
                    help_window.protocol("WM_DELETE_WINDOW", help_window.destroy)
                    help_window.wait_window()
        
        cv2.destroyAllWindows()
        sys.exit()

    def save_annotations(self, output_path):
        """保存当前标注到YOLO格式的文本文件"""
        try:
            with open(output_path, 'w') as f:
                for ann in self.annotations:
                    # 计算归一化中心坐标和宽高
                    img_h, img_w = self.original_image.shape[:2]
                    x_center = ((ann['x1'] + ann['x2']) / 2) / img_w
                    y_center = ((ann['y1'] + ann['y2']) / 2) / img_h
                    width = (ann['x2'] - ann['x1']) / img_w
                    height = (ann['y2'] - ann['y1']) / img_h
                    
                    # 写入YOLO格式: class x_center y_center width height
                    f.write(f"{ann['class']} {x_center:.6f} {y_center:.6f} {width:.6f} {height:.6f}\n")
            return True
        except Exception as e:
            print(f"保存标注时出错: {e}" if self.language == "zh" else f"Error saving annotations: {e}")
            return False

    def draw_rectangle(self, event, x, y, flags, param):
        # 转换鼠标坐标到原始图像坐标系（修复缩放偏移问题）
        if self.original_image is not None and self.image is not None:
            # 计算鼠标在显示图像中的相对位置
            rel_x = x / self.window_width
            rel_y = y / self.window_height
            
            # 计算鼠标在缩放后图像中的位置
            scaled_width = int(self.window_width * self.scale_factor)
            scaled_height = int(self.window_height * self.scale_factor)
            scaled_x = rel_x * scaled_width
            scaled_y = rel_y * scaled_height
            
            # 转换为原始图像坐标
            x = int((scaled_x + self.offset_x) / self.scale_x)
            y = int((scaled_y + self.offset_y) / self.scale_y)
        
        # 更新转换后的鼠标位置
        self.mouse_x, self.mouse_y = x, y
        
        if event == cv2.EVENT_MOUSEWHEEL:
            # 计算缩放前后的鼠标位置相对于图像的比例
            old_scale = self.scale_factor
            if flags > 0:  # 向上滚动，放大
                self.scale_factor = min(self.scale_factor * 1.1, self.max_scale)
            else:  # 向下滚动，缩小
                self.scale_factor = max(self.scale_factor * 0.9, self.min_scale)
            
            # 计算缩放中心偏移
            if old_scale != self.scale_factor:
                rel_x = (x - self.offset_x) / (self.window_width * old_scale)
                rel_y = (y - self.offset_y) / (self.window_height * old_scale)
                self.offset_x = x - rel_x * self.window_width * self.scale_factor
                self.offset_y = y - rel_y * self.window_height * self.scale_factor
        elif event == cv2.EVENT_LBUTTONDOWN:
            self.drawing = True
            self.ix, self.iy = x, y
            self.fx, self.fy = x, y  # 初始化结束坐标为起始坐标
        elif event == cv2.EVENT_MOUSEMOVE:
            if self.drawing:
                self.fx, self.fy = x, y
            # 检测鼠标悬停在哪个框上（使用转换后的坐标）
            self.hover_index = -1
            for i, ann in enumerate(self.annotations):
                if (ann['x1'] <= x <= ann['x2']) and (ann['y1'] <= y <= ann['y2']):
                    self.hover_index = i
                    break
        elif event == cv2.EVENT_LBUTTONUP:
            if self.drawing:
                self.drawing = False
                self.fx, self.fy = x, y
                # 只有当实际拖动了一定距离(>10像素)才添加标注
                if abs(self.ix - self.fx) > 10 and abs(self.iy - self.fy) > 10:
                    self.annotations.append({
                        'class': self.current_class,
                        'x1': min(self.ix, self.fx),
                        'y1': min(self.iy, self.fy),
                        'x2': max(self.ix, self.fx),
                        'y2': max(self.iy, self.fy)
                    })
                    # 自动弹出标签编辑窗口
                    self.hover_index = len(self.annotations) - 1  # 设置为最新添加的框
                    from tkinter import simpledialog
                    root = Tk()
                    root.withdraw()
                    new_label = simpledialog.askstring("编辑标签", 
                                                     "输入新标签:", 
                                                     initialvalue=self.class_names[self.current_class])
                    if new_label:  # 如果用户输入了内容
                        self.class_names[self.current_class] = new_label
                self.ix, self.iy = -1, -1
                self.fx, self.fy = -1, -1
                self.hover_index = -1  # 重置悬停索引
        elif event == cv2.EVENT_RBUTTONDOWN:  # 右键单击删除当前悬停的框
            if self.hover_index != -1:
                del self.annotations[self.hover_index]
                self.hover_index = -1
        elif event == cv2.EVENT_LBUTTONDBLCLK:  # 左键双击编辑标签
            if self.hover_index != -1:
                # 使用Tkinter输入框获取新标签
                from tkinter import simpledialog
                root = Tk()
                root.withdraw()
                new_label = simpledialog.askstring("编辑标签", 
                                                 "输入新标签:", 
                                                 initialvalue=self.class_names[self.annotations[self.hover_index]['class']])
                if new_label:  # 如果用户输入了内容
                    self.class_names[self.annotations[self.hover_index]['class']] = new_label

if __name__ == "__main__":
    # 隐藏终端窗口
    if sys.platform == "win32":
        import ctypes
        ctypes.windll.user32.ShowWindow(ctypes.windll.kernel32.GetConsoleWindow(), 0)
    
    root = Tk()
    root.withdraw()  # 隐藏Tk主窗口
    
    # 使用弹窗选择模式
    from tkinter import simpledialog
    choice = simpledialog.askstring("选择处理模式", 
                                  "请选择处理模式:\n1 - 处理整个文件夹\n2 - 处理单张图片", 
                                  initialvalue="1")
    
    if choice == '1':
        folder_path = filedialog.askdirectory(title="选择图片文件夹")
        if not folder_path:
            sys.exit()
            
        tool = YOLOLabelTool()
        valid_extensions = ['.jpg', '.jpeg', '.png', '.bmp']
        
        for filename in os.listdir(folder_path):
            if any(filename.lower().endswith(ext) for ext in valid_extensions):
                image_path = os.path.join(folder_path, filename)
                output_path = os.path.splitext(image_path)[0] + ".txt"
                
                if tool.load_image(image_path):
                    tool.run(image_path, output_path)
                    tool.annotations = []  # 清空标注为下一张图准备
                    
    elif choice == '2':
        image_path = filedialog.askopenfilename(
            title="选择图片文件",
            filetypes=[("Image files", "*.jpg *.jpeg *.png *.bmp")]
        )
        if not image_path:
            sys.exit()
            
        tool = YOLOLabelTool()
        output_path = os.path.splitext(image_path)[0] + ".txt"
        if tool.load_image(image_path):
            tool.run(image_path, output_path)
    else:
        sys.exit()

    # 添加中文字体路径（需要确保系统中有这个字体）
    # 删除这行重复定义
    # self.font_path = "C:/Windows/Fonts/simhei.ttf"  # 黑体字